import torchvision
from torch.utils import data
from torchvision import transforms

def load_data_fashion_mnist(batch_size, resize=None):  #@save
    """下载Fashion-MNIST数据集，然后将其加载到内存中"""
    # 初始化变换列表， ToTensor变换用于将PIL图像转换为张量
    trans = [transforms.ToTensor()]
    if resize:
        # 若存在图片大小调整，则在表头添加图片大小变换
        trans.insert(0, transforms.Resize(resize))
    # 将多个变换组合成一个整体变换操作
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(
        root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)
    return (data.DataLoader(mnist_train, batch_size, shuffle=True,
                            num_workers=get_dataloader_workers()),
            data.DataLoader(mnist_test, batch_size, shuffle=False,
                            num_workers=get_dataloader_workers()))

def get_dataloader_workers():  #@save
    """使用4个进程来读取数据"""
    return 4


if __name__ == '__main__':
    train_iter, test_iter = load_data_fashion_mnist(32, resize=64)
    print('begin')
    for X, y in train_iter:
        print(X.shape, X.dtype, y.shape, y.dtype)
        break
    print('done')